package com.example.util;

import com.example.config.BusinessConfig;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.utils.CloseableUtils;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import java.util.function.Consumer;

/**
 * @author: James.H.Fu
 * @description:
 * @date: Created in 21:56 2018-05-29
 * @modified By:
 */
@Component
@Slf4j
public class ZkUtil {
    @Autowired
    private BusinessConfig businessConfig;

    /**
     * zk链接
     */
    private CuratorFramework client;

    @PreDestroy
    public void onDestroy() {
        if (client != null) {
            log.info("正在销毁Bean, 尝试关闭ZooKeeper CuratorFramework链接。");
            CloseableUtils.closeQuietly(client);
        }
    }

    /**
     * 获取zk客户端
     *
     * @return
     */
    public synchronized CuratorFramework getCuratorFramework() {
        if (client == null) {
            RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 10);
            client = CuratorFrameworkFactory.builder()
                    .retryPolicy(retryPolicy)
                    .connectString(businessConfig.getZkAddress())
                    .sessionTimeoutMs(businessConfig.getZkSessionTimeout())
                    .namespace(businessConfig.getZkNamespace())
                    .build();
            client.getConnectionStateListenable().addListener((client, newState) -> {
                if (newState.isConnected()) {
                    log.info("ZooKeeper链接成功。");
                } else {
                    log.error("ZooKeeper链接失败。");
                    CloseableUtils.closeQuietly(client);
                    client = null;
                }
            });
            client.getUnhandledErrorListenable().addListener((message, throwable) -> {
                log.error("ZooKeeper出现未处理错误, message: {}", message, throwable);
            });
            client.start();
        }
        return client;
    }

    /**
     * 创建结节
     *
     * @param path
     * @param data
     * @param statConsumer
     */
    @SneakyThrows
    public void createNode(String path, String data, Consumer<Stat> statConsumer) {
        CuratorFramework client = getCuratorFramework();
        Stat stat = new Stat();
        client.create().creatingParentsIfNeeded().storingStatIn(stat).forPath(path, data == null ? "".getBytes() : data.getBytes());
        if (statConsumer != null) {
            statConsumer.accept(stat);
        }
    }

    /**
     * 创建节点
     *
     * @param path
     * @param data
     * @throws Exception
     */
    public void createNode(String path, String data) throws Exception {
        CuratorFramework client = getCuratorFramework();
        Stat s = client.checkExists().forPath(path);
        if (s == null) {
            client.create().creatingParentsIfNeeded().forPath(path, data == null ? "".getBytes() : data.getBytes());
        }
    }

    /**
     * 创建节点
     *
     * @param path
     * @param createMode
     * @param data
     * @throws Exception
     */
    public void createNode(String path, String data, CreateMode createMode) throws Exception {
        CuratorFramework client = getCuratorFramework();
        Stat s = client.checkExists().forPath(path);
        if (s == null) {
            client.create().withMode(createMode).forPath(path, data == null ? "".getBytes() : data.getBytes());
        }
    }

    /**
     * 获取数据
     *
     * @param path
     * @return
     * @throws Exception
     */
    public String getData(String path) throws Exception {
        CuratorFramework client = getCuratorFramework();
        return new String(client.getData().forPath(path));
    }

    /**
     * 更新数据
     *
     * @param path
     * @param data
     * @throws Exception
     */
    public void updateData(String path, String data) throws Exception {
        CuratorFramework client = getCuratorFramework();
        client.setData().forPath(path, data.getBytes());

    }

    /**
     * 删除节点
     *
     * @param path
     * @throws Exception
     */
    public void deleteNode(String path) throws Exception {
        CuratorFramework client = getCuratorFramework();
        client.delete().forPath(path);
    }
}
